import tkinter as tk
from tkinter import filedialog, messagebox
from PIL import Image, ImageTk
import tensorflow as tf
import numpy as np

class TrafficSignApp:
    def __init__(self, root):
        self.root = root
        self.root.title("交通标志识别系统")
        self.root.geometry("800x600")
        
        # 加载训练好的模型
        try:
            self.model = tf.keras.models.load_model('traffic_sign_model.h5')
        except:
            messagebox.showerror("错误", "无法加载模型文件，请确保traffic_sign_model.h5存在")
            self.root.destroy()
            return
        
        # 创建GUI组件
        self.create_widgets()
        
        # 类别标签（示例，实际应该与训练时的类别对应）
        self.class_names = [
            "限速5", "限速15", "限速30", "限速40", "限速50", 
            "限速60", "限速70","限速80","禁止超车", "禁止货车通行", "禁止停车", 
            "禁止鸣笛", "注意行人", "注意儿童", "注意自行车",
            "注意野生动物", "前方施工", "危险", "左转", "右转",
            "直行", "环岛", "让行", "停车", "禁止车辆通行",
            "其他标志"  
        ]
    
    def create_widgets(self):
        # 顶部标题
        title_frame = tk.Frame(self.root)
        title_frame.pack(pady=10)
        
        tk.Label(
            title_frame, 
            text="交通标志识别系统", 
            font=("Arial", 16, "bold")
        ).pack()
        
        # 图像显示区域
        self.image_frame = tk.Frame(self.root, bg="white", width=500, height=400)
        self.image_frame.pack(pady=20)
        self.image_frame.pack_propagate(False)
        
        self.image_label = tk.Label(self.image_frame)
        self.image_label.pack(expand=True, fill="both")
        
        # 结果显示区域
        self.result_label = tk.Label(
            self.root, 
            text="识别结果将显示在这里", 
            font=("Arial", 12),
            fg="blue"
        )
        self.result_label.pack(pady=10)
        
        # 置信度显示
        self.confidence_label = tk.Label(
            self.root, 
            text="置信度: 0%", 
            font=("Arial", 10)
        )
        self.confidence_label.pack()
        
        # 按钮区域
        button_frame = tk.Frame(self.root)
        button_frame.pack(pady=20)
        
        tk.Button(
            button_frame, 
            text="选择图片", 
            command=self.load_image,
            width=15,
            height=2
        ).pack(side=tk.LEFT, padx=10)
        
        tk.Button(
            button_frame, 
            text="识别标志", 
            command=self.predict_image,
            width=15,
            height=2
        ).pack(side=tk.LEFT, padx=10)
        
        tk.Button(
            button_frame, 
            text="退出系统", 
            command=self.root.quit,
            width=15,
            height=2
        ).pack(side=tk.LEFT, padx=10)
    
    def load_image(self):
        file_path = filedialog.askopenfilename(
            filetypes=[("Image files", "*.jpg *.jpeg *.png")]
        )
        if file_path:
            try:
                self.current_image = Image.open(file_path)
                self.display_image(self.current_image)
                self.result_label.config(text="图片已加载", fg="green")
            except Exception as e:
                messagebox.showerror("错误", f"无法加载图片: {str(e)}")
    
    def display_image(self, image):
        # 调整图像大小以适应显示区域
        img_width, img_height = image.size
        frame_width = self.image_frame.winfo_width()
        frame_height = self.image_frame.winfo_height()
        
        ratio = min(frame_width/img_width, frame_height/img_height)
        new_size = (int(img_width*ratio), int(img_height*ratio))
        resized_image = image.resize(new_size, Image.LANCZOS)
        
        self.tk_image = ImageTk.PhotoImage(resized_image)
        self.image_label.config(image=self.tk_image)
    
    def predict_image(self):
        if not hasattr(self, 'current_image'):
            messagebox.showwarning("警告", "请先选择一张图片")
            return
        
        try:
            # 预处理图像以匹配模型输入
            img = self.current_image.resize((128, 128))
            img_array = np.array(img) / 255.0
            img_array = np.expand_dims(img_array, axis=0)  # 添加批次维度
            
            # 进行预测
            predictions = self.model.predict(img_array)
            predicted_class = np.argmax(predictions[0])
            confidence = np.max(predictions[0]) * 100
            
            # 显示结果
            if predicted_class < len(self.class_names):
                result_text = f"识别结果: {self.class_names[predicted_class]}"
            else:
                result_text = "识别结果: 未知标志"
            
            self.result_label.config(text=result_text, fg="blue")
            self.confidence_label.config(text=f"置信度: {confidence:.2f}%")
            
        except Exception as e:
            messagebox.showerror("错误", f"识别过程中出错: {str(e)}")

if __name__ == "__main__":
    root = tk.Tk()
    app = TrafficSignApp(root)
    root.mainloop()